package com.boot2.core.web.xss;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
* @author zhangweilin
* @date 2020/10/14 18:28
* @Description:  可以自定义添加header的request
*/
public final class ModifiableHeaderHttpServletRequest extends HttpServletRequestWrapper {
    private final Map<String, String> customHeaderMap;

    public ModifiableHeaderHttpServletRequest(HttpServletRequest request){
        super(request);
        this.customHeaderMap = new HashMap<>();
    }

    /**
     * 自定义添加header
     * @param name
     * @param value
     */
    public void putHeader(String name, String value){
        this.customHeaderMap.put(name, value);
    }

    @Override
    public String getHeader(String name) {
        String headerValue = customHeaderMap.get(name);

        if (headerValue != null){
            return headerValue;
        }
        return ((HttpServletRequest) getRequest()).getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        Set<String> set = new HashSet<String>(customHeaderMap.keySet());

        @SuppressWarnings("unchecked")
        Enumeration<String> e = ((HttpServletRequest) getRequest()).getHeaderNames();
        while (e.hasMoreElements()) {
            String n = e.nextElement();
            set.add(n);
        }

        return Collections.enumeration(set);
    }
}